#comments-start
Copyright (c) <2011>, <John Seekins>
All rights reserved.

Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
    * Neither the name of the <ORGANIZATION> nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#comments-end

#cs
Problems with bad disk detection appeared to be related to diskpart processes not being closed properly.  They are now.
#ce

#include <Constants.au3>
#include <Array.au3>

Func BasicDiskData()
	Local $preDskList = ""
	Local $procID = _DiskProcessStart()

	StdinWrite($procID, "LIST DISK" & @LF)
	Sleep(10)
	$preDskList = _ReadFromProcess($procID, "###")
	_DiskProcessEnd($procID)

	If StringInStr($preDskList, "2003") <> 0 Then
		Local $startOfList = StringInStr($preDskList, "###")
		$preDskList = StringTrimLeft($preDskList, $startOfList)
	EndIf
	Local $startOfList = StringInStr($preDskList, "0")

	If $startOfList = 0 Then
		MsgBox(64, "Warning!", "No fixed disks are attached to this computer!")
		Return 0
	EndIf

	Local $endOfList = StringInStr($preDskList, "B", 1, -1)
	Local $diskList = StringTrimRight($preDskList, StringLen($preDskList) - $endOfList)
	$diskList = StringTrimLeft($diskList, $startOfList - 6)
	Local $diskArrayString = StringSplit($diskList, @LF)

	Local $diskArray[$diskArrayString[0] + 1][2]
	$diskArray[0][0] = $diskArrayString[0]
	For $i = 1 To $diskArrayString[0]
		$diskArrayString[$i] = StringStripWS($diskArrayString[$i], 8)
		Local $startOfListSize = StringInStr($diskArrayString[$i], "e", 1)
		$diskArray[$i][0] = StringTrimLeft($diskArrayString[$i], $startOfListSize)
		Local $startOfListFree = StringInStr($diskArray[$i][0], "B", 1)
		Local $endOfListSize = StringInStr($diskArray[$i][0], "B", 1)
		$diskArray[$i][1] = StringTrimLeft($diskArray[$i][0], $startOfListFree)
		$diskArray[$i][0] = StringTrimRight($diskArray[$i][0], StringLen($diskArray[$i][0]) - $endOfListSize)
		$diskArray[$i][0] = _ConvertToMB($diskArray[$i][0])
		$diskArray[$i][1] = _ConvertToMB($diskArray[$i][1])
	Next
	Return $diskArray
EndFunc

Func DetailedDiskData($diskSizeArray, $diskNumber)
	Local $partNum = _FindPartitions($diskSizeArray, $diskNumber)
	If $partNum = 0 Then
		Return 0
	Else
		Local $partListArray[$partNum + 1][3]
		$partListArray[0][0] = $partNum

		$procID = _DiskProcessStart()
		StdinWrite($procID, "SELECT DISK " & $diskNumber - 1 & @LF)
		For $i = 1 To $partNum
			StdinWrite($procID, "SELECT PARTITION " & $i & @LF)
			StdinWrite($procID, "DETAIL PARTITION" & @LF)

			$dataString = _ReadFromProcess($procID, "Active")

			Local $trimPoint = StringInStr($dataString, "Active")
			$dataString = StringTrimLeft($dataString, $trimPoint + 6)
			$activeString = StringTrimRight($dataString, StringLen($dataString) - 4)
			$partListArray[$i][0] = StringStripWS($activeString, 8)

			$trimPoint = StringInStr($dataString, "NTFS")
			If $trimPoint = 0 Then
				$trimPoint = StringInStr($dataString, "FAT32")
			EndIf
			$dataString = StringTrimLeft($dataString, $trimPoint - 1)
			$activeString = StringTrimRight($dataString, StringLen($dataString) - 5)
			$partListArray[$i][1] = StringStripWS($activeString, 8)

			$trimPoint = StringInStr($dataString, "B", 1)
			$activeString = StringTrimRight($dataString, $trimPoint + 5)
			$activeString = StringTrimLeft($activeString, 16)
			$activeString = StringStripWS($activeString, 8)
			$partListArray[$i][2] = _ConvertToMB($activeString)
		Next
		_DiskProcessEnd($procID)
		return $partListArray
	EndIf
EndFunc

Func CreatePartition($simpleArray, $diskNumber, $partSize)
	Local $procID = _DiskProcessStart()

	$partSize = _ConvertToMB($partSize)

	If ($simpleArray[$diskNumber][1] < $partSize) OR ($simpleArray[$diskNumber][0] < $partSize) Then
		MsgBox(64, "Warning", "Cannot create a partition that exceeds the free or total space on the disk... ")
		Return 0
	EndIf
	StdinWrite($procID, "SELECT DISK " & $diskNumber - 1 & @LF)
	StdinWrite($procID, "CREATE PARTITION PRIMARY SIZE=" & $partSize & @LF)
	;StdinWrite($procID, "
	If Not @error Then
		MsgBox(64, "Success!", "Successfully created " & $partSize & "MB partition on disk " & $diskNumber)
		_DiskProcessEnd($procID)
		Return 1
	Else
		MsgBox(64, "Error!", "Error creating partition... ")
		_DiskProcessEnd($procID)
		Return 0
	EndIf
EndFunc

Func DeletePartition($diskNumber, $partNum)
	Local $procID = _DiskProcessStart()

	StdinWrite($procID, "SELECT DISK " & $diskNumber - 1 & @LF)
	StdinWrite($procID, "SELECT PARTITION " & $partNum & @LF)
	StdinWrite($procID, "DELETE PARTITION" & @LF)
	If Not @error Then
		MsgBox(64, "Success!", "Successfully deleted partition " & $partNum & " on disk " & $diskNumber)
		_DiskProcessEnd($procID)
		Return 1
	Else
		MsgBox(64, "Error!", "Error deleting partition!")
		_DiskProcessEnd($procID)
		Return 0
	EndIf
EndFunc

Func SetPartitionActive($diskNumber, $partNum)
	Local $procID = _DiskProcessStart()

	StdinWrite($procID, "SELECT DISK " & $diskNumber - 1 & @LF)
	StdinWrite($procID, "SELECT PARTITION " & $partNum & @LF)
	StdinWrite($procID, "ACTIVE" & @LF)

	If Not @error Then
		MsgBox(64, "Success!", "Successfully set partition " & $partNum & " of disk " & $diskNumber & " to active!")
		_DiskProcessEnd($procID)
		Return 1
	Else
		MsgBox(64, "Error!", "Error setting active partition")
		_DiskProcessEnd($procID)
		Return 0
	EndIf
EndFunc

Func _DiskProcessStart()
	Return Run("DiskPart.exe", "", @SW_HIDE, $STDIN_CHILD + $STDOUT_CHILD)
EndFunc

Func _ReadFromProcess($procID, $breakString)
;Sleep(100)
	While 1
		sleep(100)
		$read = StdoutRead($procID)
		;if the particular string "breakstring" (which is case-sensitive) is found in the output, jump out of the loop
		If StringInStr($read, $breakString, 2) <> 0 Then ExitLoop
	Wend
	Return $read
EndFunc

Func _ConvertToMB($diskSize)
	If StringLen($diskSize) > 2 Then
		Local $sizeType = StringTrimLeft($diskSize, StringLen($diskSize) - 2)
		Local $MBSize = Number(StringTrimRight($diskSize, 2))

		Switch $sizeType
			Case "KB"
				Return $MBSize / 1000
			Case "MB"
				Return $MBSize
			Case "GB"
				Return $MBSize * 1000
			Case "TB"
				Return $MBSize * 1000 * 1000
		EndSwitch
	Else
		Return 0
	EndIf
EndFunc

Func _DiskProcessEnd($procID)
	;StdinWrite($procID)
	StdioClose($procID)
	ProcessClose($procID)
EndFunc

Func _FindPartitions($diskSizeArray, $diskNumber)
	Local $dataString = ""
	Local $procID = _DiskProcessStart()

	StdinWrite($procID, "SELECT DISK " & $diskNumber - 1 & @LF)
	StdinWrite($procID, "LIST PARTITION" & @LF)

	$dataString = _ReadFromProcess($procID, "partition")

	$lastPartPos = StringInStr($dataString, "Partition", 1, -1)
	If $lastPartPos <> 0 Then
		Local $partNum = StringTrimLeft($dataString, $lastPartPos + 9)
		$partNum = StringTrimRight($partNum, StringLen($partNum) - 1)
		$partNum = StringStripWS($partNum, 8)

		_DiskProcessEnd($procID)
		Return $partNum
	Else
		_DiskProcessEnd($procID)
		Return 0
	EndIf
EndFunc
